'''
Created on May 31, 2010

@author: Guan Gui, Will Zhang
$LastChangedBy: guiguan $
'''
from collections import defaultdict
import math

__version__ = "$Rev: 359 $"

class LanguageModel(object):
    """
    This class implements a N-Gram Language Model for Phrase Based Translation 
    Systems. 
    """

    def __init__(self, n, aligned_sent_collection):
        '''
        Initialize a new C{LanguageModel} using n-gram model
        
        @param n: n-gram size
        @type n: C{int}
        @param aligned_sent_collection: collection of all C{AlignedSent} in a 
        given corpus
        @type aligned_sent_collection: C{AlignedSentCollection}
        '''
        self._ngram = NGramModel(n, aligned_sent_collection)
        self._n = n
        if n > 1:
            self._nminus1gram = NGramModel(n - 1, aligned_sent_collection)
        else:
            self._nminus1gram = None

    @property
    def n(self):
        '''
        The order of ngram
        '''
        return self._n

    def __getitem__(self, key):
        '''
        Query the language model
        
        @param key: a tuple of the form (e1, e2, e3...)
        @type key: C{tuple}
        @return: P(e1 | e2 e3 ...)
        @rtype: float
        '''
        assert len(key) <= self._n, \
        "You cannot query LM with higher order than N:%d" % (self._n)
        if self._n == 1:
            return self._ngram[key]
        elif self._n == 2 and len(key) == 1:
            return self._nminus1gram[key]
        elif len(key) == self._n:
            return self._ngram[key] / self._nminus1gram[key[0:-1]]
        else:
            return 1

class NGramModel(defaultdict):

    def __init__(self, n, aligned_sent_collect):
        '''
        Initialize a new C{NGramModel} with order n
        
        @param n: n-gram size
        @type n: C{int}
        @param aligned_sent_collect: collection of all C{AlignedSent} in a 
        given corpus
        @type aligned_sent_collect: C{AlignedSentCollection}
        '''
        base = super(NGramModel, self)
        base.__init__(float)
        self._n = n
        self._default_value = 0
        num_ngram = 0
        for aligned_sent in aligned_sent_collect.sents:
            for i in xrange(0, len(aligned_sent.words) - n):
                key = tuple(aligned_sent.words[i:i + n])
                base.__setitem__(key, base.__getitem__(key) + 1)
                num_ngram += 1
        b = len(aligned_sent_collect.set_words)
        b = math.pow(b, n)
        denominator = num_ngram + b
        self._default_value = 1.0 / denominator
        for k in self:
            base.__setitem__(k, (1.0 + base.__getitem__(k)) / denominator)

    def __getitem__(self, key):
        '''
        Query the ngram model
        
        @param key: a tuple of the form (e1, e2, e3...)
        @type key: C{tuple}
        @return: P(e1 e2 e3 ...)
        @rtype: float
        '''
        if key in self:
            return super(NGramModel, self).__getitem__(key)
        else:
            return self._default_value

    @property
    def n(self):
        '''
        The order of ngram
        '''
        return self._n

if __name__ == '__main__':
    pass
